#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Sep 29 02:36:52 2020

@author: sarupa
"""

import mpctools as mpc
import control
import numpy as np
import matplotlib.pyplot as plt
import time
from numpy import random
from fullsystem_9states_EKF import *
from random import seed
import random
from casadi import *
from numpy import linalg as la

start = time.perf_counter()

#Define RMSE
def rmse(predictions, targets):

    differences = (predictions - targets)/targets                      

    differences_squared = differences ** 2                    

    mean_of_differences_squared = differences_squared.mean()  

    rmse_val = np.sqrt(mean_of_differences_squared)*100         

    return rmse_val 

# Fix random seed
np.random.seed(100) # Seed random number generator. 1, 5

# Load data
x_test_unscaled=np.loadtxt('x_test_unscaled.txt')#to get the inputs 
n= len(x_test_unscaled)
u_test=x_test_unscaled[0:n,4:10]

# Problem parameters.

Delta = 0.01
print('\n step size',Delta)
Nsim = len(u_test[:, 1])-2 #Number of data points
tplot = np.arange(Nsim+1)*Delta

Nx = 9  # Numbeof system states
Nu = 6   # Number of system inputsn
Ny = 3 # Number of system outputs
Nw = Nx  # Number of process disturbances
Nv = Ny  # Number of measurement noise


# Make covariance matrices. # Covariance for prior.

sigma_v = 0.005 #Standard deviation of the measurements
sigma_w = 100 #Standard deviation for the process noise
sigma_p = 20 #Standard deviation for prior
#Best tuning parameter found for the full estimator

P = np.diag((sigma_p*np.ones((Nx,)))**2) 
Q = np.diag((sigma_w*np.ones((Nw,)))**2)
R = np.diag((sigma_v*np.ones((Nv,)))**2)

# Import data
x0=np.array([9.06393325e-01, 9.32535771e-02, 3.81482083e+02, 5.80306393e-01,
          3.36753698e-01, 3.68003376e+02, 3.18882689e-01, 6.21370214e-01,
        3.76773315e+02]) #actual
x_0 =x0*0.9 #initial guess


# Define the casadi variables
x_symbol = SX.sym('x',Nx)
u_symbol=SX.sym('u',Nu)
w_symbol=SX.sym('w',Nw)

# Make a simulator. MPC simulator
model_cstr_casadi = mpc.DiscreteSimulator(cstr_xy_ode_scale, Delta, [Nx,Nu,Nw], ["x","u","w"],casaditype='SX')    
print("Hi, I did the discretization")
# Convert continuous-time f to explicit discrete-time F with RK4.
F = mpc.getCasadiFunc(cstr_xy_ode_scale,[Nx,Nu,Nw],["x","u","w"],"F",rk4=True,Delta=Delta,M=1)
H = mpc.getCasadiFunc(measurement_xy_model,[Nx],["x"],"H")

#Define the measurement noise
sigma_v_r=0.005
v = sigma_v_r*np.random.randn(Nsim,Nv)

sigma_w_r=0.00
w = sigma_w_r*np.random.randn(Nsim,Nw)

# Set the input vector to the cstr
u=u_test
usim=np.zeros((Nsim, Nu))
usim[:,:] = u[:-2,:] # We use the input vector to the process instead of dummy input

# Define variables
xsim = np.zeros((Nsim+1,Nx))
xsim[0,:] = x0
yclean = np.zeros((Nsim, Ny))
ysim = np.zeros((Nsim, Ny))

# Simulate the process dynamics
for t in range(Nsim):
    yclean[t,:] = measurement_xy_model(xsim[t]) # Get zero-noise measurement.
    ysim[t,:] = yclean[t,:] + v[t,:] # Add noise to measurement.#xsim[t,8]+v[t,:]#    
    xsim[t+1,:] = model_cstr_casadi.sim(xsim[t,:],usim[t,:],w[t,:])   


# Now do estimation.
xhat_ = np.zeros((Nsim+1,Nx))
xhat = np.zeros((Nsim,Nx))
yhat = np.zeros((Nsim,Ny))
vhat = np.zeros((Nsim,Nv))
what = np.zeros((Nsim,Nw))
x0bar = x_0
xhat[0,:] = x0bar
guess = {}
Pmp1 = P

start2 = time.perf_counter()
for t in range(Nsim-1):
#     # Define sizes of everything.           
#     # Apply model function to get xhat(t+1 | t )
#     #xhat_[t+1,:] = np.squeeze(F(xhat[t,:], usim[t,:], np.zeros((Nw,))))
# #============================================================================#            
#         # Do EKF to update prior covariance, but don't take EKF state. Remove if arrival cost is not needed
    [Pmp1, xhatmp1, Pkk, xkk,A] = mpc.ekf(F,H,x=xhat[t,:],u=usim[t,:],w=w[t,:],y=ysim[t,:],P=Pmp1,Q=Q,R=R)
    xhat[t,:] = xkk
    xhat[t+1,:]= xhatmp1

finish2 = time.perf_counter()
print('\n Finished online simulation in time',finish2 - start2,'seconds')


# Open loop simulation without feedback
xsimo = np.zeros((Nsim+1,Nx))
xsimo[0,:] = x_0
ycleano = np.zeros((Nsim, Ny))
ysimo = np.zeros((Nsim, Ny))
for t in range(Nsim):
    ycleano[t,:] = measurement_xy_model(xsimo[t]) # Get zero-noise measurement.
    ysimo[t,:] = ycleano[t,:] + v[t,:] # Add noise to measurement.    
    xsimo[t+1,:] = model_cstr_casadi.sim(xsimo[t,:],usim[t,:],w[t,:])


#error calculation
error=abs(xsim[-1]-xhat)

str_output = ['CB3']
p1=1
for i in range(p1): 
    
    plt.subplot(p1,1,i+1)
    plt.plot(xhat[:Nsim,7],'b')
    plt.plot(xsim[:Nsim,7],'m')
    plt.plot(x_test_unscaled[:Nsim,0],':g')
    plt.ylabel("Temp " + str(i+1))
    plt.ylabel(str_output[i])
    plt.legend(['EKF','Actual'])
    plt.grid()
    plt.xlabel("Time")
    if i == 0:
        plt.title('9 states EKF vs Actual vs Open loop')
plt.legend(['EKF','Actual', 'Open'])

#define C matrix
C = np.array([[0.,0.,1.,0.,0.,0.,0.,0.,0],
                         [0.,0.,0.,0.,0.,1.,0.,0.,0],
                         [0.,0.,0.,0.,0.,0.,0.,0.,1.]]) 
# degree of observability
O_A=control.obsv(A,C)
rank_O=la.matrix_rank(O_A)
print('Rank of a matrix',rank_O)
print('conditioning no=' , np.linalg.cond(O_A))
print('degree of observability=' , 1/np.linalg.cond(O_A))

# RMSE is reported after taking average of using different initial guess and random seed
#this is a sample calculation

rmse_cb_ekf=rmse(xhat[:400,7],xsim[0:400,7])
print('rmse_cb_ekf',rmse_cb_ekf)

rmse_cb_open=rmse(xsimo[:Nsim,7],xsim[0:Nsim,7])
print('rmse_cb_open',rmse_cb_open)

finish = time.perf_counter()
print('\n Finished entire simulation in',finish - start,'seconds')
